from manim import *
import multiprocessing

COMPARE_POS1 = [-5, 3, 0]
COMPARE_POS = [-4, 3, 0]
COMPARE_POS2 = [-3, 3, 0]
GAP_POS = [-5, 3, 0]
SIDE_LENGTH = 1
POSX = -6
POSY = -3
POSZ = 0


class Element:
    value:int
    mob:any

    def __init__(self, value, mob):
        self.value = value
        self.mob = mob

class Sort(Scene):

    def construct(self):
        ValueArr = [5, 9, 1, 7, 12, 3, 10, 2, 8, 6, 4, 11]
        # ValueArr = [5, 9, 1, 7, 3, 10]
        Arr = []
        i = 0
        while i < len(ValueArr):
            element_shape = Rectangle(width = SIDE_LENGTH,
                                      height = SIDE_LENGTH,
                                      color = "#808080",
                                      fill_color = WHITE,
                                      stroke_width=4,
                                      fill_opacity=1)
            element_text = Text(str(ValueArr[i]), font_size = 48, color = BLACK)
            element = VGroup(element_shape, element_text)
            element.move_to([POSX + i * SIDE_LENGTH, POSY, POSZ])
            Arr.append(Element(ValueArr[i], element))
            self.add(element)
            self.bring_to_front(element_text)
            self.wait(0.1)
            i += 1

        # self.InsertSort(Arr, N)
        self.ShellSort(Arr)

    def Move(self, e, POS):
        self.play(e.mob.animate.move_to(POS))

    def Swap(self, e1, e2):
        POS1 = e1.mob.get_center()
        POS2 = e2.mob.get_center()
        self.play(e1.mob.animate.move_to(POS2),
                  e2.mob.animate.move_to(POS1))

    def Compare(self, e1, e2, swap):
        animations = []
        e1.mob[1].set_z_index(2)
        e2.mob[1].set_z_index(2)
        animations.append(e1.mob[0].animate.set_fill(RED,opacity=1).set_stroke(e1.mob[0].stroke_color))
        animations.append(e2.mob[0].animate.set_fill(RED,opacity=1).set_stroke(e2.mob[0].stroke_color))
        self.play(AnimationGroup(*animations, lag_ratio=0))
        self.wait()

        string = str(e1.value)
        if e1.value > e2.value:
            string += " > "
        elif e1.value < e2.value:
            string += " < "
        else:
            string += " = "
        string += str(e2.value)

        if swap:
            string += ", 交换"
        else:
            string += ", 不交换"

        text = Text(string, color = RED)
        text.move_to(COMPARE_POS)
        self.play(Create(text))
        self.wait(0.5)
        self.play(FadeOut(text))



        animations = []
        e1.mob[1].set_z_index(2)
        e2.mob[1].set_z_index(2)
        animations.append(e1.mob[0].animate.set_fill(GREEN,opacity=1).set_stroke(e1.mob[0].stroke_color))
        animations.append(e2.mob[0].animate.set_fill(GREEN,opacity=1).set_stroke(e2.mob[0].stroke_color))
        self.play(AnimationGroup(*animations, lag_ratio=0))
        self.wait()

    def InsertSort(self, Arr, n, gap = 1):
        i = 1

        self.play(Arr[0].mob[0].animate.set_fill(GREEN,opacity=1).set_stroke(Arr[0].mob[0].stroke_color))
        while i < n:
            temp = Arr[i]
            temp_pos = temp.mob.get_center()
            self.play(temp.mob.animate.shift(UP * 2))
            j = i - 1
            while j >= 0:
                if Arr[j].value > temp.value:
                    self.Compare(Arr[j], temp, 1)
                    temp_pos = Arr[j].mob.get_center()
                    self.play(Arr[j].mob.animate.shift([0 + SIDE_LENGTH * gap, 0, 0]),
                              temp.mob.animate.shift([0 - SIDE_LENGTH * gap, 0, 0]))
                    Arr[j + 1] = Arr[j]
                elif Arr[j].value < temp.value:
                    self.Compare(Arr[j], temp, 0)
                    break
                else:
                    self.Compare(Arr[j], temp, 0)
                    break
                j -= 1

            Arr[j + 1] = temp
            self.Move(temp, temp_pos)
            i += 1

    def ShellSort(self, Arr):
        gap = len(Arr) // 2

        while gap > 0:
            
            nSubArr = 0
            while nSubArr < gap:
                SubArr = []
                indices = []
                i = nSubArr   
                while i < len(Arr):
                    SubArr.append(Arr[i])
                    indices.append(i)
                    i += gap

                animations = []
                i = 0
                while i < len(SubArr):
                    vg = SubArr[i].mob
                    vg[1].set_z_index(2)
                    animations.append(vg[0].animate.set_fill(BLUE, opacity=1).set_stroke(vg[0].stroke_color)) 
                    i += 1

                self.play(AnimationGroup(*animations, lag_ratio=0))
                self.wait()

                self.InsertSort(SubArr, len(SubArr), gap)
                
                for idx, elem in zip(indices, SubArr):
                    Arr[idx] = elem

                animations = []
                i = 0
                while i < len(SubArr):
                    vg = SubArr[i].mob
                    vg[1].set_z_index(2)
                    animations.append(vg[0].animate.set_fill(WHITE, opacity=1).set_stroke(vg[0].stroke_color)) 
                    i += 1

                self.play(AnimationGroup(*animations, lag_ratio=0))
                self.wait()

                nSubArr += 1

            gap //= 2

if __name__ == "__main__":
    config.max_cores = multiprocessing.cpu_count()  # 使用所有核心
    config.quality = "high_quality"
    scene = Sort()
    scene.render()